
from abc import abstractmethod
from pev2.data.DataModule import DataModule
from pandas import DataFrame
from matplotlib import pyplot as plt

class BaseStockPool:

    '''
    股票池的初始化方法
    :param startDate: 开始日期
    :param endDate: 结束日期
    :param interval: 再平衡周期
    '''
    def __init__(self, startDate, endDate, interval):
        self.startDate = startDate
        self.endDate = endDate
        self.interval = interval

    """
    获取股票池的出票列表
    :return:
    """
    @abstractmethod
    def getOptionStocks(self):
        pass

    """
    统计股票池的收益
    """
    def statisticProfit(self):
        # 设定评测周期
        rebalanceDates, codesDict = self.getOptionStocks()

        dm = DataModule()

        # 用DataFrame保存收益
        dfProfit = DataFrame(columns=['profit', 'hs300'])
        dfProfit.loc[rebalanceDates[0]] = {'profit': 0, 'hs300': 0}

        # 获取沪深300在统计周期内的第一天的值
        hs300K = dm.getKData('000300', index=True, startDate=rebalanceDates[0], endDate=rebalanceDates[0])
        hs300BeginValue = hs300K.loc[hs300K.index[0]]['close']

        # 通过净值计算累计收益
        netValue = 1
        for index in range(1, len(rebalanceDates) - 1):
            lastRebalanceDate = rebalanceDates[index - 1]
            currentRebalanceDate = rebalanceDates[index]

            # 获取上一期的股票池
            codes = codesDict[lastRebalanceDate]

            # 统计当前的收益
            profitSum = 0

            # 参与统计收益的股票个数
            profitCodeCount = 0

            for code in codes:
                dailyKS = dm.getKData(code, autype='hfq', startDate=lastRebalanceDate, endDate=currentRebalanceDate)
                indexSize = dailyKS.index.size

                # 如果没有数据，则跳过，长期停牌
                if indexSize == 0:
                    continue
                # 买入价
                inPrice = dailyKS.loc[dailyKS.index[0]]['close']
                # 卖出价
                outPrice = dailyKS.loc[dailyKS.index[indexSize - 1]]['close']
                # 股票池内所有的股票收益
                profitSum += (outPrice - inPrice) / inPrice
                profitCodeCount += 1

            profit = round(profitSum / profitCodeCount, 4)
            hs300KCurrent = dm.getKData('000300', index=True, startDate=currentRebalanceDate, endDate=currentRebalanceDate)
            hs300Close = hs300KCurrent.loc[hs300KCurrent.index[0]]['close']

            # 计算净值和累计收益
            netValue = netValue * (1 + profit)
            dfProfit.loc[currentRebalanceDate] = {
                'profit': round((netValue - 1) * 100, 4),
                'hs300': round((hs300Close - hs300BeginValue) * 100 / hs300BeginValue, 4)
            }
            print(dfProfit)

        # 绘制曲线
        dfProfit.plot(title='stock pool statistic', kind='line')
        # 显示图像
        plt.show()
